{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/NeuromatchAcademy/course-content/blob/master/tutorials/W3D2_HiddenDynamics/W3D2_Tutorial3.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Neuromatch Academy: Week 3, Day 2, Tutorial 3\n",
    "# Linear Dynamical Systems & The Kalman Filter in 1D\n",
    "__Content creators:__ Itzel Olivos Castillo and Xaq Pitkow\n",
    "\n",
    "__Content reviewers:__ Your names here\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Tutorial Objectives\n",
    "\n",
    "In previous tutorials we used Hidden Markov Models (HMM) to infer *discrete* latent states from a sequence of measurements. In this tutorial, we will learn how to infer a *continuous* latent variable using the Kalman filter, which is one version of an HMM.\n",
    "\n",
    "In this tutorial, you will:\n",
    "* Review linear dynamical systems\n",
    "* Learn about the Kalman filter in one dimension\n",
    "* Manipulate parameters of process to see how the Kalman filter behaves\n",
    "* Think about some core properties of the Kalman filter.\n",
    "\n",
    "You can imagine this inference process happening as Mission Control tries to locate and track Astrocat. But you can also imagine that the brain is using an analogous Hidden Markov Model to track objects in the world, or to estimate the consequences of its own actions. And you could use this technique to estimate brain activity from noisy measurements, for understanding or for building a brain-machine interface."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-03T14:03:11.524104Z",
     "iopub.status.busy": "2021-06-03T14:03:11.523437Z",
     "iopub.status.idle": "2021-06-03T14:03:12.046875Z",
     "shell.execute_reply": "2021-06-03T14:03:12.045836Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import transforms\n",
    "from collections import namedtuple\n",
    "from scipy.stats import norm\n",
    "\n",
    "gaussian = namedtuple('Gaussian', ['mean', 'cov'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:03:12.051930Z",
     "iopub.status.busy": "2021-06-03T14:03:12.051312Z",
     "iopub.status.idle": "2021-06-03T14:03:12.207342Z",
     "shell.execute_reply": "2021-06-03T14:03:12.206741Z"
    }
   },
   "outputs": [],
   "source": [
    "# @title Figure Settings\n",
    "import ipywidgets as widgets       # interactive display\n",
    "from ipywidgets import interactive, interact, HBox, Layout,VBox\n",
    "from IPython.display import HTML\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "plt.style.use(\"https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/nma.mplstyle\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:03:12.219021Z",
     "iopub.status.busy": "2021-06-03T14:03:12.218422Z",
     "iopub.status.idle": "2021-06-03T14:03:12.222063Z",
     "shell.execute_reply": "2021-06-03T14:03:12.222484Z"
    }
   },
   "outputs": [],
   "source": [
    "# @title Plotting Functions\n",
    "\n",
    "def plot_trajectory(s, T):\n",
    "  fig, ax = plt.subplots()\n",
    "  ax.plot([0, T], [0, 0], 'k')\n",
    "  ax.plot(s, color='limegreen', lw=2)\n",
    "  ax.plot([T], [s[-1]], marker='o', markersize=8, color='limegreen')\n",
    "  ax.set(xlabel = 'Time t', ylabel = 's(t)')\n",
    "\n",
    "def plot_measurement(s, m, T):\n",
    "  fig, axes = plt.subplots(1, 2)\n",
    "\n",
    "  axes[0].plot(s, color='limegreen', lw=2, label='true position')\n",
    "  axes[0].plot([T], [s[-1]], marker='o', markersize=8, color='limegreen')\n",
    "  axes[0].plot(m, '.', color='crimson', lw=2, label='measurement')\n",
    "  axes[0].set(xlabel = 'Time t', ylabel = 's(t)')\n",
    "\n",
    "  sbounds = 1.1*max(max(np.abs(s)), max(np.abs(m)))\n",
    "  axes[1].plot([-sbounds, sbounds], [-sbounds, sbounds], 'k')    # plot line of equality\n",
    "\n",
    "  axes[1].scatter(s, m, marker='.', color='red', s=100)\n",
    "  axes[1].set(xlabel = 'State', ylabel = 'Measurement')\n",
    "\n",
    "def plot_function(u=1,v=2,w=3,x=4,y=5,z=6):\n",
    "    time=np.arange(0,1,0.01)\n",
    "    df=pd.DataFrame({\"Y1\":np.sin(time*u*2*np.pi),\"y2\":np.sin(time*v*2*np.pi),\"y3\":np.sin(time*w*2*np.pi),\n",
    "                    \"y4\":np.sin(time*x*2*np.pi),\"y5\":np.sin(time*y*2*np.pi),\"y6\":np.sin(time*z*2*np.pi)})\n",
    "    df.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:03:12.232942Z",
     "iopub.status.busy": "2021-06-03T14:03:12.231434Z",
     "iopub.status.idle": "2021-06-03T14:03:12.233894Z",
     "shell.execute_reply": "2021-06-03T14:03:12.234368Z"
    }
   },
   "outputs": [],
   "source": [
    "# @title Helper Functions\n",
    "\n",
    "# Fixed params\n",
    "T_max = 200\n",
    "D = 1\n",
    "tau_min = 1\n",
    "tau_max = 50\n",
    "process_noise_min = 0.1\n",
    "process_noise_max = 10\n",
    "measurement_noise_min = 0.1\n",
    "measurement_noise_max = 10\n",
    "\n",
    "unit_process_noise = np.random.randn(T_max)     # compute all N(0, 1) in advance to speed up time slider\n",
    "unit_measurement_noise = np.random.randn(T_max)     # compute all N(0, 1) in advance to speed up time slider\n",
    "\n",
    "def filter(D, process_noise, measurement_noise, belief, m, s):\n",
    "    prior = gaussian(belief.mean, belief.cov)\n",
    "\n",
    "    predicted_estimate = D * belief.mean\n",
    "    predicted_covariance = D**2 * belief.cov + process_noise\n",
    "    yesterdays_post = gaussian(predicted_estimate, predicted_covariance)\n",
    "\n",
    "    likelihood = gaussian(m, measurement_noise)\n",
    "    innovation_estimate = m - predicted_estimate\n",
    "    innovation_covariance = predicted_covariance + measurement_noise\n",
    "\n",
    "    K = predicted_covariance / innovation_covariance  # Kalman gain, i.e. the weight given to the difference between the measurement and predicted measurement\n",
    "    updated_mean = predicted_estimate + K * innovation_estimate\n",
    "    updated_cov = (1 - K) * predicted_covariance\n",
    "    todays_posterior = gaussian(updated_mean, updated_cov)\n",
    "\n",
    "    return prior, yesterdays_post, likelihood, todays_posterior"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Section 1: Astrocat Dynamics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:03:12.242583Z",
     "iopub.status.busy": "2021-06-03T14:03:12.241927Z",
     "iopub.status.idle": "2021-06-03T14:03:12.304400Z",
     "shell.execute_reply": "2021-06-03T14:03:12.304868Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Video 1: Astrocat through time\n",
    "# Insert the ID of the corresponding youtube video\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"P1jFKXDzAYQ\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:03:12.310927Z",
     "iopub.status.busy": "2021-06-03T14:03:12.310306Z",
     "iopub.status.idle": "2021-06-03T14:03:12.348170Z",
     "shell.execute_reply": "2021-06-03T14:03:12.347601Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Video 2: Introduction\n",
    "# Insert the ID of the corresponding youtube video\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"hqYP3-M3rFg\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 1.1: Simulating Astrocat's movements\n",
    "\n",
    "Astrocat's position $s_t$ is determined by the linear dynamical system $s_t = Ds_{t-1} + w_{t-1}$. $D$ is a scalar that models how astrocat would like to change its position over time and $w_t \\sim \\mathcal{N}(0, \\sigma_p^2)$ is white Gaussian noise caused by unreliable actuators in astrocat's propulsion unit."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Coding Exercise 1.1: Simulating Astrocat's movements\n",
    "\n",
    "We will first simulate Astrocat's movements according to the linear dynamical system outlined above. \n",
    "\n",
    "You will complete a function `simulate_trajectory` below to simulate a possible trajectory,\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-03T14:03:12.357227Z",
     "iopub.status.busy": "2021-06-03T14:03:12.353923Z",
     "iopub.status.idle": "2021-06-03T14:03:12.454469Z",
     "shell.execute_reply": "2021-06-03T14:03:12.453910Z"
    }
   },
   "outputs": [],
   "source": [
    "def simulate_trajectory(D, s0, sigma_eta, T):\n",
    "  \"\"\" Compute the response of the linear dynamical system.\n",
    "\n",
    "  Args:\n",
    "    D (scalar): dynamics multiplier\n",
    "    s0 (scalar): initial postion\n",
    "    sigma_eta (scalar): amount of noise in the system\n",
    "    T (scalar): total duration of the simulation\n",
    "\n",
    "  Returns:\n",
    "    ndarray: `s`: astrocat's trajectory up to time T\n",
    "  \"\"\"\n",
    "\n",
    "  ###################################################################\n",
    "  ## Fill out the following then remove\n",
    "  raise NotImplementedError(\"Student exercise: need to implement simulation\")\n",
    "  ###################################################################\n",
    "\n",
    "  # Initialize variables\n",
    "  s = np.zeros(T + 1)\n",
    "  s[0] = s0\n",
    "\n",
    "  # Compute the position at time t given the position at time t-1 for all t\n",
    "  # Consider that np.random.normal(mu, sigma) generates a random sample from\n",
    "  # a gaussian with mean = mu and standard deviation = sigma\n",
    "\n",
    "  for t in range(1, len(s)):\n",
    "\n",
    "    s[t] = ...\n",
    "\n",
    "  return s\n",
    "\n",
    "# Set random seed\n",
    "np.random.seed(0)\n",
    "\n",
    "# Choose parameters\n",
    "D = -0.5    # parameter in f(x)\n",
    "T = 10      # total Time duration\n",
    "s0 = 5.     # initial condition of x at time 0\n",
    "sigma_eta = 2 # amount of noise in the actuators of astrocat's propulsion unit\n",
    "\n",
    "# Simulate trajectory\n",
    "s = simulate_trajectory(D, s0, sigma_eta, T)\n",
    "\n",
    "# Visualize\n",
    "plot_trajectory(s, T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-03T14:03:12.461187Z",
     "iopub.status.busy": "2021-06-03T14:03:12.460597Z",
     "iopub.status.idle": "2021-06-03T14:03:12.742800Z",
     "shell.execute_reply": "2021-06-03T14:03:12.742168Z"
    }
   },
   "outputs": [],
   "source": [
    "# to_remove solution\n",
    "def simulate_trajectory(D, s0, sigma_eta, T):\n",
    "  \"\"\" Compute the response of the linear dynamical system.\n",
    "\n",
    "  Args:\n",
    "    D (scalar): dynamics multiplier\n",
    "    s0 (scalar): initial postion\n",
    "    sigma_eta (scalar): amount of noise in the system\n",
    "    T (scalar): total duration of the simulation\n",
    "\n",
    "  Returns:\n",
    "    ndarray: `s`: astrocat's trajectory up to time T\n",
    "  \"\"\"\n",
    "\n",
    "  # Initialize variables\n",
    "  s = np.zeros(T + 1)\n",
    "  s[0] = s0\n",
    "\n",
    "  # Compute the position at time t given the position at time t-1 for all t\n",
    "  # Consider that np.random.normal(mu, sigma) generates a random sample from\n",
    "  # a gaussian with mean = mu and standard deviation = sigma\n",
    "\n",
    "  for t in range(1, len(s)):\n",
    "\n",
    "    s[t] = D*s[t-1] + np.random.normal(0, sigma_eta)\n",
    "\n",
    "  return s\n",
    "\n",
    "# Set random seed\n",
    "np.random.seed(0)\n",
    "\n",
    "# Choose parameters\n",
    "D = -0.5    # parameter in f(x)\n",
    "T = 10      # total Time duration\n",
    "s0 = 5.     # initial condition of x at time 0\n",
    "sigma_eta = 2 # amount of noise in the actuators of astrocat's propulsion unit\n",
    "\n",
    "# Simulate trajectory\n",
    "s = simulate_trajectory(D, s0, sigma_eta, T)\n",
    "\n",
    "# Visualize\n",
    "with plt.xkcd():\n",
    "  plot_trajectory(s, T)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interactive Demo 1.1: Astrocat movement\n",
    "\n",
    "In the demo below, you can play with the parameters, specially $D$ and the noise standard deviation $\\sigma_p$, of the linear dynamical system governing Astrocat's movement. We will look at a longer time span.\n",
    "\n",
    "\n",
    "\n",
    "1.   What happens when D is large (>1)? Why?\n",
    "2.   What happens when D is a large negative number (<-1)? Why?\n",
    "3.   What about when D is zero?\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:03:12.775758Z",
     "iopub.status.busy": "2021-06-03T14:03:12.750799Z",
     "iopub.status.idle": "2021-06-03T14:03:13.031788Z",
     "shell.execute_reply": "2021-06-03T14:03:13.031263Z"
    }
   },
   "outputs": [],
   "source": [
    "# @markdown Execute this cell to enable the widget\n",
    "\n",
    "def simulate_trajectory(D, s0, sigma_eta, T):\n",
    "  \"\"\" Compute the response of the linear dynamical system.\n",
    "\n",
    "  Args:\n",
    "    D (scalar): dynamics multiplier\n",
    "    s0 (scalar): initial postion\n",
    "    sigma_eta (scalar): amount of noise in the system\n",
    "    T (scalar): total duration of the simulation\n",
    "\n",
    "  Returns:\n",
    "    ndarray: `s`: astrocat's trajectory up to time T\n",
    "  \"\"\"\n",
    "\n",
    "  # Initialize variables\n",
    "  s = np.zeros(T + 1)\n",
    "  s[0] = s0\n",
    "\n",
    "  # Compute the position at time t given the position at time t-1 for all t\n",
    "  # Consider that np.random.normal(mu, sigma) generates a random sample from\n",
    "  # a gaussian with mean = mu and standard deviation = sigma\n",
    "\n",
    "  for t in range(1, len(s)):\n",
    "\n",
    "    s[t] = D*s[t-1] + np.random.normal(0, sigma_eta)\n",
    "\n",
    "  return s\n",
    "\n",
    "@widgets.interact(D = widgets.FloatSlider(value=-.5, min=-2, max=2, step=0.1), sigma_p = widgets.FloatSlider(value=2.0, min=0, max=5, step=0.1))\n",
    "\n",
    "def interactive_movement(D, sigma_p):\n",
    "\n",
    "    # Choose parameters\n",
    "    T = 100      # total Time duration\n",
    "    s0 = 5.     # initial condition of x at time 0\n",
    "\n",
    "    # Simulate trajectory\n",
    "    s = simulate_trajectory(D, s0, sigma_p, T)\n",
    "\n",
    "    # Visualize\n",
    "    plot_trajectory(s, T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-03T14:03:13.039207Z",
     "iopub.status.busy": "2021-06-03T14:03:13.038524Z",
     "iopub.status.idle": "2021-06-03T14:03:13.041671Z",
     "shell.execute_reply": "2021-06-03T14:03:13.042185Z"
    }
   },
   "outputs": [],
   "source": [
    "# to_remove explanation\n",
    "\n",
    "\"\"\"\n",
    "\n",
    "1) When D is large, the state at time step t will depend heavily on the state at time\n",
    "   step t_1. If we forget about the noise term, D = 2 would mean that the state at each\n",
    "   time step is double the one before! So the state becomes huge and basically explodes towards\n",
    "   infinity.\n",
    "\n",
    "2) If D is a large negative number, the state at time t will be a different sign than the\n",
    "   state at time step t_1. So the state will oscillate over the x axis.\n",
    "\n",
    "3) When D is zero, the state at time t will not depend on the previous state, it will just\n",
    "   be drawn from the noise distribution\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 1.2: Measuring Astrocat's movements"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Coding Exercise 1.2: Reading measurements from Astrocat's collar\n",
    "\n",
    "Since we will not have access to Astrocat's true position, we must estimate it using measurements of a noisy sensor attached to its collar. Complete the function below to read measurements from Astrocat's collar.\n",
    "\n",
    "We show the measurement and state at every time step in the left plot. In the right plot, we make a scatter plot to compare measured state to true state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-03T14:03:13.181440Z",
     "iopub.status.busy": "2021-06-03T14:03:13.157855Z",
     "iopub.status.idle": "2021-06-03T14:03:13.451850Z",
     "shell.execute_reply": "2021-06-03T14:03:13.452317Z"
    }
   },
   "outputs": [],
   "source": [
    "def read_collar(s, sigma_measurements):\n",
    "  \"\"\" Compute the measurements of the noisy sensor attached to astrocat's collar\n",
    "\n",
    "  Args:\n",
    "    s (ndarray): astrocat's true position over time\n",
    "    sigma_measurements (scalar): amount of noise in the sensor\n",
    "\n",
    "  Returns:\n",
    "    ndarray: `m`: astrocat's position over time according to the sensor\n",
    "  \"\"\"\n",
    "\n",
    "  # Initialize variables\n",
    "  m = np.zeros(len(s))\n",
    "\n",
    "  # For all time t, add white Gaussian noise with magnitude sigma_measurements\n",
    "  # Consider that np.random.normal(mu, sigma) generates a random sample from\n",
    "  # a gaussian with mean = mu and standard deviation = sigma\n",
    "\n",
    "  for t in range(len(s)):\n",
    "\n",
    "    ###################################################################\n",
    "    ## Fill out the following then remove\n",
    "    ## raise NotImplementedError(\"Student exercise: need to implement read_collar function\")\n",
    "    ###################################################################\n",
    "\n",
    "    m[t] = s[t] + sigma_measurements * np.random.randn()\n",
    "\n",
    "  return m\n",
    "\n",
    "# Choose parameters\n",
    "sigma_measurement = 4 # amount of noise in astrocat's collar\n",
    "\n",
    "# Read measurements from color\n",
    "m = read_collar(s, sigma_measurement)\n",
    "\n",
    "# Visualize\n",
    "plot_measurement(s, m, T)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interactive Demo 1.2: Compare true states to measured states\n",
    "\n",
    "In the following demo, you can change $\\sigma_{measurement}$ and visualize the resulting measurements and measurement accuracy.\n",
    "\n",
    "Before playing with the demo, discuss:\n",
    "\n",
    "\n",
    "1.   What is the difference between $\\sigma_{measurement}$ and the $\\sigma_p$ we played with in Interactive Demo 1.1?\n",
    "2. How will accuracy of measured state change with $\\sigma_{measurement}$? \n",
    "3.   What will happen if $\\sigma_{measurement}$ is 0?\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:03:13.474727Z",
     "iopub.status.busy": "2021-06-03T14:03:13.474183Z",
     "iopub.status.idle": "2021-06-03T14:03:13.856327Z",
     "shell.execute_reply": "2021-06-03T14:03:13.857159Z"
    }
   },
   "outputs": [],
   "source": [
    "# @markdown Execute this cell to enable the widget\n",
    "\n",
    "@widgets.interact(sigma_m = widgets.FloatSlider(value=2.0, min=0, max=5, step=0.1))\n",
    "def plot_measurement_quality(sigma_m):\n",
    "\n",
    "  # Read measurements from color\n",
    "  m = read_collar(s, sigma_m)\n",
    "\n",
    "  # Visualize\n",
    "  plot_measurement(s, m, T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-03T14:03:13.865608Z",
     "iopub.status.busy": "2021-06-03T14:03:13.864992Z",
     "iopub.status.idle": "2021-06-03T14:03:13.871690Z",
     "shell.execute_reply": "2021-06-03T14:03:13.870473Z"
    }
   },
   "outputs": [],
   "source": [
    "# to_remove explanation\n",
    "\n",
    "\"\"\"\n",
    "\n",
    "1) sigma_m is how much noise there is in the measurement observation model\n",
    "   (the reading of the color location).\n",
    "   sigma_p was the noise in the actual movement every step.\n",
    "\n",
    "2) Accuracy will be worse with higher noise\n",
    "\n",
    "3) If sigma_m is 0, the measured states will be fully accurate.\n",
    "\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see by the measured state quality in the previous demo, using only the measurements to track astrocat can be catastrophic. A Kalman filter will solve this problem, as we will see in the next section!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Section 2: The Kalman filter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:03:13.877680Z",
     "iopub.status.busy": "2021-06-03T14:03:13.877019Z",
     "iopub.status.idle": "2021-06-03T14:03:13.914736Z",
     "shell.execute_reply": "2021-06-03T14:03:13.915167Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Video 3: Explanation of Kalman filter\n",
    "# Insert the ID of the corresponding youtube video\n",
    "from IPython.display import YouTubeVideo\n",
    "video = YouTubeVideo(id=\"EFCS55MtB-s\", width=854, height=480, fs=1)\n",
    "print(\"Video available at https://youtu.be/\" + video.id)\n",
    "video"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 2.1: Using the Kalman filter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interactive Demo 2.1: The Kalman filter in action\n",
    "\n",
    "Next we provide you with an interactive visualization to understand how the Kalman filter works. Play with the sliders to gain an intuition for how the different factors affect the Kalman filter's inferences.\n",
    "\n",
    "The sliders:\n",
    "* current time: Kalman filter synthesizes measurements up until this time.\n",
    "* dynamics time constant $\\tau$: this determines the dynamics matrix, $D=\\exp{-\\Delta t/\\tau}$ where $\\Delta t$ is the discrete time step (here 1).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:03:13.922897Z",
     "iopub.status.busy": "2021-06-03T14:03:13.917628Z",
     "iopub.status.idle": "2021-06-03T14:03:14.290632Z",
     "shell.execute_reply": "2021-06-03T14:03:14.291130Z"
    }
   },
   "outputs": [],
   "source": [
    "#@markdown Execute this cell to enable the widget\n",
    "display(HTML('''<style>.widget-label { min-width: 15ex !important; }</style>'''))\n",
    "\n",
    "@widgets.interact(T=widgets.IntSlider(T_max/4, description=\"current time\", min=2, max=T_max-1),\n",
    "                  tau=widgets.FloatSlider(tau_max/2, description='dynamics time constant', min=tau_min, max=tau_max),\n",
    "                  process_noise=widgets.FloatSlider(2, description=\"process noise\", min=process_noise_min, max=process_noise_max),\n",
    "                  measurement_noise=widgets.FloatSlider(3, description=\"observation noise\", min=measurement_noise_min, max=measurement_noise_max),\n",
    "                  flag_s = widgets.Checkbox(value=True, description='state', disabled=True, indent=False),\n",
    "                  flag_m = widgets.Checkbox(value=False, description='measurement', disabled=False, indent=False),\n",
    "                  flag_s_ = widgets.Checkbox(value=False, description='estimate', disabled=False, indent=False),\n",
    "                  flag_err_ = widgets.Checkbox(value=False, description='estimator confidence intervals', disabled=False, indent=False))\n",
    "\n",
    "def stochastic_system(T, tau, process_noise, measurement_noise, flag_m, flag_s_, flag_err_):\n",
    "    t = np.arange(0, T_max, 1)              # timeline\n",
    "    s = np.zeros(T_max)                     # states\n",
    "    D = np.exp(-1/tau)                      # dynamics multiplier (matrix if s is vector)\n",
    "    prior_mean = 0\n",
    "    prior_cov = process_noise/(1-D**2)\n",
    "\n",
    "    s[0] = np.sqrt(prior_cov) * unit_process_noise[0]   # Sample initial condition from equilibrium distribution\n",
    "\n",
    "    m = np.zeros(T_max)    # measurement\n",
    "    s_ = np.zeros(T_max)    # estimate (posterior mean)\n",
    "    cov_ = np.zeros(T_max)    # uncertainty (posterior covariance)\n",
    "\n",
    "    s_[0] = prior_mean\n",
    "    cov_[0] = prior_cov\n",
    "    posterior = gaussian(prior_mean, prior_cov)\n",
    "\n",
    "    captured_prior = None\n",
    "    captured_likelihood = None\n",
    "    captured_posterior = None\n",
    "\n",
    "    for i in range(1, T_max):\n",
    "        s[i] = D * s[i-1] + np.sqrt(process_noise) * unit_process_noise[i-1]\n",
    "        m[i] = s[i] + np.sqrt(measurement_noise) * unit_measurement_noise[i]\n",
    "\n",
    "        prior, yesterdays_post, likelihood, posterior = filter(D, process_noise, measurement_noise, posterior, m[i], s[i])\n",
    "\n",
    "        s_[i] =  posterior.mean\n",
    "        cov_[i] = posterior.cov\n",
    "\n",
    "        if i == T:\n",
    "          captured_prior = prior\n",
    "          captured_likelihood = likelihood\n",
    "          captured_posterior = posterior\n",
    "\n",
    "    smin = min(min(m),min(s-2*np.sqrt(cov_[-1])),min(s_-2*np.sqrt(cov_[-1])))\n",
    "    smax = max(max(m),max(s+2*np.sqrt(cov_[-1])),max(s_+2*np.sqrt(cov_[-1])))\n",
    "    pscale = 0.2  # scaling factor for displaying pdfs\n",
    "\n",
    "    fig = plt.figure(figsize=[15, 10])\n",
    "    ax = plt.subplot(2, 1, 1)\n",
    "    ax.set_xlabel('time')\n",
    "    ax.set_ylabel('state')\n",
    "    ax.set_xlim([0, T_max+(T_max*pscale)])\n",
    "    ax.set_ylim([smin, smax])\n",
    "\n",
    "    show_pdf = [False, False]\n",
    "    ax.plot(t[:T+1], s[:T+1], color='limegreen', lw=2)\n",
    "    ax.plot(t[T:], s[T:], color='limegreen', lw=2, alpha=0.3)\n",
    "    ax.plot([t[T:T+1]], [s[T:T+1]], marker='o', markersize=8, color='limegreen')\n",
    "\n",
    "    if flag_m:\n",
    "        ax.plot(t[:T+1], m[:T+1], '.', color='crimson', lw=2)\n",
    "        ax.plot([t[T:T+1]], [m[T:T+1]], marker='o', markersize=8, color='crimson')\n",
    "\n",
    "        domain = np.linspace(ax.get_ylim()[0], ax.get_ylim()[1], 500)\n",
    "        pdf_likelihood = norm.pdf(domain, captured_likelihood.mean, np.sqrt(captured_likelihood.cov))\n",
    "        ax.fill_betweenx(domain, T + pdf_likelihood*(T_max*pscale), T, color='crimson', alpha=0.5, label='likelihood', edgecolor=\"crimson\", linewidth=0)\n",
    "        ax.plot(T + pdf_likelihood*(T_max*pscale), domain, color='crimson', linewidth=2.0)\n",
    "\n",
    "        ax.legend(ncol=3, loc='upper left')\n",
    "        show_pdf[0] = True\n",
    "\n",
    "    if flag_s_:\n",
    "        ax.plot(t[:T+1], s_[:T+1], color='black', lw=2)\n",
    "        ax.plot([t[T:T+1]], [s_[T:T+1]], marker='o', markersize=8, color='black')\n",
    "        show_pdf[1] = True\n",
    "\n",
    "    if flag_err_:\n",
    "        ax.fill_between(t[:T+1], s_[:T+1] + 2 * np.sqrt(cov_)[:T+1], s_[:T+1] - 2 * np.sqrt(cov_)[:T+1], color='black', alpha=0.3)\n",
    "        show_pdf[1] = True\n",
    "\n",
    "    if show_pdf[1]:\n",
    "        domain = np.linspace(ax.get_ylim()[0], ax.get_ylim()[1], 500)\n",
    "        pdf_post = norm.pdf(domain, captured_posterior.mean, np.sqrt(captured_posterior.cov))\n",
    "        ax.fill_betweenx(domain, T + pdf_post*(T_max*pscale), T, color='black', alpha=0.5, label='posterior', edgecolor=\"black\", linewidth=0)\n",
    "        ax.plot(T + pdf_post*(T_max*pscale), domain, color='black', linewidth=2.0)\n",
    "        ax.legend(ncol=3, loc='upper left')\n",
    "\n",
    "    if show_pdf[0] and show_pdf[1]:\n",
    "        domain = np.linspace(ax.get_ylim()[0], ax.get_ylim()[1], 500)\n",
    "        pdf_prior = norm.pdf(domain, captured_prior.mean, np.sqrt(captured_prior.cov))\n",
    "        ax.fill_betweenx(domain, T + pdf_prior*(T_max*pscale), T, color='dodgerblue', alpha=0.5, label='prior', edgecolor=\"dodgerblue\", linewidth=0)\n",
    "        ax.plot(T + pdf_prior*(T_max*pscale), domain, color='dodgerblue', linewidth=2.0)\n",
    "        ax.legend(ncol=3, loc='upper left')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 2.2: Estimation Accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interactive Demo 2.2: Compare states, estimates, and measurements\n",
    "\n",
    "How well do the estimates $\\hat{s}$ match the actual values $s$? How does the distribution of errors $\\hat{s}_t - s_t$ compare to the posterior variance? Why? Try different parameters of the Hidden Markov Model and observe how the properties change.\n",
    "\n",
    "How do the _measurements_ $m$ compare to the true states?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:03:14.308748Z",
     "iopub.status.busy": "2021-06-03T14:03:14.307979Z",
     "iopub.status.idle": "2021-06-03T14:03:14.999044Z",
     "shell.execute_reply": "2021-06-03T14:03:14.998477Z"
    }
   },
   "outputs": [],
   "source": [
    "#@markdown Execute this cell to enable the widget\n",
    "display(HTML('''<style>.widget-label { min-width: 15ex !important; }</style>'''))\n",
    "\n",
    "@widgets.interact(tau=widgets.FloatSlider(tau_max/2, description='tau', min=tau_min, max=tau_max),\n",
    "                  process_noise=widgets.FloatSlider(2, description=\"process noise\", min=process_noise_min, max=process_noise_max),\n",
    "                  measurement_noise=widgets.FloatSlider(3, description=\"observation noise\", min=measurement_noise_min, max=measurement_noise_max),\n",
    "                  flag_m = widgets.Checkbox(value=False, description='measurements', disabled=False, indent=False))\n",
    "\n",
    "def stochastic_system(tau, process_noise, measurement_noise, flag_m):\n",
    "    T = T_max\n",
    "    t = np.arange(0, T_max, 1)              # timeline\n",
    "    s = np.zeros(T_max)                     # states\n",
    "    D = np.exp(-1/tau)                      # dynamics multiplier (matrix if s is vector)\n",
    "    prior_mean = 0\n",
    "    prior_cov = process_noise/(1-D**2)\n",
    "\n",
    "    s[0] = np.sqrt(prior_cov) * np.random.randn()   # Sample initial condition from equilibrium distribution\n",
    "\n",
    "    m = np.zeros(T_max)    # measurement\n",
    "    s_ = np.zeros(T_max)    # estimate (posterior mean)\n",
    "    cov_ = np.zeros(T_max)    # uncertainty (posterior covariance)\n",
    "\n",
    "    s_[0] = prior_mean\n",
    "    cov_[0] = prior_cov\n",
    "    posterior = gaussian(prior_mean, prior_cov)\n",
    "\n",
    "    for i in range(1, T):\n",
    "        s[i] = D * s[i-1] + np.sqrt(process_noise) * np.random.randn()\n",
    "        m[i] = s[i] + np.sqrt(measurement_noise) * np.random.randn()\n",
    "\n",
    "        prior, yesterdays_post, likelihood, posterior = filter(D, process_noise, measurement_noise, posterior, m[i], s[i])\n",
    "\n",
    "        s_[i] =  posterior.mean\n",
    "        cov_[i] = posterior.cov\n",
    "\n",
    "    fig = plt.figure(figsize=[10, 5])\n",
    "    ax = plt.subplot(1, 2, 1)\n",
    "    ax.set_xlabel('s')\n",
    "    ax.set_ylabel('$\\mu$')\n",
    "\n",
    "    sbounds = 1.1*max(max(np.abs(s)), max(np.abs(s_)), max(np.abs(m)))\n",
    "    ax.plot([-sbounds, sbounds], [-sbounds, sbounds], 'k')    # plot line of equality\n",
    "    ax.errorbar(s, s_, yerr=2*np.sqrt(cov_[-1]), marker='.', mfc='black', mec='black', linestyle='none', color='gray')\n",
    "\n",
    "    axhist = plt.subplot(1, 2, 2)\n",
    "    axhist.set_xlabel('error $s-\\hat{s}$')\n",
    "    axhist.set_ylabel('probability')\n",
    "    axhist.hist(s-s_, density=True, bins=25, alpha=.5, label='histogram of estimate errors', color='yellow')\n",
    "\n",
    "    if flag_m:\n",
    "        ax.plot(s, m, marker='.', linestyle='none', color='red')\n",
    "        axhist.hist(s-m,density=True,bins=25,alpha=.5,label='histogram of measurement errors',color='orange')\n",
    "\n",
    "    domain = np.arange(-sbounds, sbounds, 0.1)\n",
    "    pdf_g = norm.pdf(domain, 0, np.sqrt(cov_[-1]))\n",
    "    axhist.fill_between(domain, pdf_g, color='black', alpha=0.5, label=r'posterior shifted to mean')\n",
    "    axhist.legend()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 2.3: How long does it take to find astrocat?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interactive Demo 2.3: Search duration\n",
    "\n",
    "Here we plot the posterior variance as a function of time. Before mission control gets measurements, their only information about astrocat's location is the prior. After some measurements, they home in on astrocat.\n",
    "* How does the variance shrink with time?\n",
    "* The speed depends on the process dynamics, but does it also depend on the signal-to-noise ratio (SNR)? (Here we measure SNR in decibels, a log scale where 1 dB means 0.1 log unit.)\n",
    "\n",
    "The red curve shows how rapidly the latent variance equilibrates from an initial condition. (It's exponential with a time constant of $\\sim 1/(1-D^2)$.) Does the latent process converge faster or slower than the posterior converges? Can you explain this based on how the Kalman filter integrates evidence?\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {
     "iopub.execute_input": "2021-06-03T14:03:15.013353Z",
     "iopub.status.busy": "2021-06-03T14:03:15.012703Z",
     "iopub.status.idle": "2021-06-03T14:03:15.308136Z",
     "shell.execute_reply": "2021-06-03T14:03:15.308638Z"
    }
   },
   "outputs": [],
   "source": [
    "#@title Initial transient for finding Astrocat\n",
    "display(HTML('''<style>.widget-label { min-width: 15ex !important; }</style>'''))\n",
    "\n",
    "@widgets.interact(T=widgets.IntSlider(tau_max/2, description=\"max time\", min=2, max=T_max-1),\n",
    "                  tau=widgets.FloatSlider(tau_max/2, description='time constant', min=tau_min, max=tau_max),\n",
    "                  SNRdB=widgets.FloatSlider(-20., description=\"SNR (decibels)\", min=-40., max=10.))\n",
    "\n",
    "def stochastic_system(T, tau, SNRdB):\n",
    "    t = np.arange(0, T, 1)              # timeline\n",
    "    s = np.zeros(T)                     # states\n",
    "    D = np.exp(-1/tau)                  # dynamics matrix (scalar here)\n",
    "    prior_mean = 0\n",
    "    process_noise = 1\n",
    "    SNR = 10**(.1*SNRdB)\n",
    "    measurement_noise = process_noise / SNR\n",
    "    prior_cov = process_noise/(1-D**2)\n",
    "\n",
    "    s[0] = np.sqrt(prior_cov) * unit_process_noise[0]   # Sample initial condition from equilibrium distribution\n",
    "\n",
    "    m = np.zeros(T)    # measurements\n",
    "    s_ = np.zeros(T)    # estimates (posterior mean)\n",
    "    cov_ = np.zeros(T)    # uncertainty (posterior covariance)\n",
    "\n",
    "    s_[0] = prior_mean\n",
    "    cov_[0] = prior_cov\n",
    "    posterior = gaussian(prior_mean, prior_cov)\n",
    "\n",
    "    for i in range(1, T):\n",
    "        s[i] = D * s[i-1] + np.sqrt(process_noise) * unit_process_noise[i-1]\n",
    "        m[i] = s[i] + np.sqrt(measurement_noise) * unit_measurement_noise[i]\n",
    "\n",
    "        prior, yesterdays_post, likelihood, posterior = filter(D, process_noise, measurement_noise, posterior, m[i], s[i])\n",
    "\n",
    "        s_[i] =  posterior.mean\n",
    "        cov_[i] = posterior.cov\n",
    "\n",
    "    equilibrium_var = process_noise * (D**2 - 1 - SNR + np.sqrt((D**2 - 1 - SNR)**2 + 4 * D**2 * SNR)) / (2 * D**2 * SNR)\n",
    "\n",
    "    fig = plt.figure()\n",
    "    plt.xlabel('time')\n",
    "    plt.ylabel('variance')\n",
    "    plt.xlim([0, T])\n",
    "    plt.ylim([0, max(cov_)])\n",
    "\n",
    "    plt.fill_between(t, 0, cov_, color='black', alpha=0.3)\n",
    "    plt.plot(t, cov_, color='black', label='posterior variance')\n",
    "    plt.plot(t, equilibrium_var + (prior_cov - equilibrium_var) * np.exp(-2*t/tau), color='red', label='autocorrelation of hidden process')\n",
    "    plt.legend()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Section 3: Applications of Kalman filter in brain science\n",
    "\n",
    "* Brain-Computer Interface: estimate intended movements using neural activity as measurements.\n",
    "* Data analysis: estimate brain activity from noisy measurements (e.g. EEG)\n",
    "* Model of perception: prey tracking using noisy sensory measurements\n",
    "* Imagine your own! When are you trying to estimate something you cannot see directly?\n",
    "\n",
    "There are many variants that improve upon the limitations of the Kalman filter: non-Gaussian states and measurements, nonlinear dynamics, and more."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Summary\n",
    "\n",
    "\n",
    "**Useful references:**\n",
    "- Roweis, Ghahramani (1998): A unifying review of linear Gaussian Models\n",
    "- Bishop (2006): Pattern Recognition and Machine Learning"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "W3D2_Tutorial3",
   "provenance": [],
   "toc_visible": true
  },
  "kernel": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
